package org.apache.pdfbox.examples.textsearch;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.pdmodel.PDPage;
import org.apache.pdfbox.pdmodel.common.PDRectangle;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.pdfbox.text.TextPosition;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * @author yuanzhenchao
 */
public class PdfTextSearch extends PDFTextStripper {
  private String text;
  private char[] textChars;
  private PDDocument document;
  private int rotation = 0;
  private List<Float[]> rectangles = new ArrayList();

  /**
   *
   * @param pdfBytes PDF文件
   * @param text 搜索的文字
   * @param rotation 文字的rotation
   * @throws IOException
   */
  public PdfTextSearch(byte[] pdfBytes, String text, int rotation)
          throws IOException {
    this.text = text;
    this.rotation = rotation;
    this.textChars = this.text.toCharArray();
    this.document = PDDocument.load(pdfBytes);
  }

  public PdfTextSearch(byte[] pdfBytes, String text)
          throws IOException {
    this.text = text;
    this.textChars = this.text.toCharArray();
    this.document = PDDocument.load(pdfBytes);
  }

  /**
   * 搜索
   * @return List&lt;Float[]{page, llx, lly, urx, ury}&gt;
   * @throws IOException
   */
  public List<Float[]> search() throws IOException {
    int numberOfPages = document.getNumberOfPages();
    for (int i = 1; i <= numberOfPages; i++) {
      this.setStartPage(i);
      this.setEndPage(i);
      PDPage page = document.getPage(i - 1);
      this.getText(document);
      Iterator<List<TextPosition>> iterator = this.charactersByArticle.iterator();
      while (iterator.hasNext()) {
        List<TextPosition> textPositions = iterator.next();
        if (!textPositions.isEmpty() && textPositions.size() >= textChars.length) {
          int offset = 0;
          while (offset <= textPositions.size() - textChars.length) {
            boolean matched = true;
            for (int j = 0; j < textChars.length; j++) {
              TextPosition textPosition = textPositions.get(offset + j);
              if (textPosition.getUnicode().charAt(0) != textChars[j] || textPosition.getDir() != rotation) {
                matched = false;
              }
            }
            if (matched) {
              List<TextPosition> macthedPositions = new ArrayList<>();
              for (int j = 0; j < textChars.length; j++) {
                TextPosition textPosition = textPositions.get(offset + j);
                macthedPositions.add(textPosition);
              }
              PDRectangle rectangle = getRectangle(macthedPositions.get(0), macthedPositions.get(macthedPositions.size() - 1), page);
              Float[] floats = new Float[]{(float)i, rectangle.getLowerLeftX(), rectangle.getLowerLeftY(), rectangle.getUpperRightX(), rectangle.getUpperRightY()};
              rectangles.add(floats);
            }
            offset++;
          }
        }
      }
    }
    return rectangles;
  }

  private PDRectangle getRectangle(TextPosition first, TextPosition end, PDPage page) {
    float lx = first.getXDirAdj();
    float ly = page.getMediaBox().getHeight() - first.getYDirAdj();

    float ux = end.getXDirAdj();
    float uy = page.getMediaBox().getHeight() - end.getYDirAdj();
    float endFontSize = end.getFontSize();
    ux = ux + endFontSize;
    uy = uy + endFontSize;

    PDRectangle rectangle = new PDRectangle();
    rectangle.setLowerLeftX(lx);
    rectangle.setLowerLeftY(ly);
    rectangle.setUpperRightX(ux);
    rectangle.setUpperRightY(uy);
    return rectangle;
  }
}
